# coding=utf-8
# Copyright (C) xxx team - All Rights Reserved
#
# @Version:   3.9.4
# @Software:  PyCharm
# @FileName:  inference.py
# @CTime:     2021/5/3 16:30
# @Author:    Haiyang Yu
# @Email:     xxx
# @UTime:     2021/5/3 16:30
#
# @Description:
#     used for inference
#     xxx
#
import os
import logging
from typing import List, Dict
import hydra
from omegaconf import DictConfig, OmegaConf
from transformers import AutoTokenizer
# self
from architecture import Architecture as Model
from metrics import get_entities
import matplotlib.pyplot as plt
import seaborn as sns

logger = logging.getLogger(__name__)


@hydra.main(config_path="conf", config_name="config")
def main(cfg: DictConfig):
    cfg.cwd = hydra.utils.get_original_cwd()
    # logger.info(OmegaConf.to_yaml(cfg))

    # load the model and dataset
    model = Model(cfg)

    model = model.load_from_checkpoint(os.path.join(cfg.cwd, cfg.load_ckpt))
    tokenizer = AutoTokenizer.from_pretrained('hfl/rbt3')
    tag_map = {0: 'O', 1: 'B-CONT', 2: 'B-ORG', 3: 'B-RACE', 4: 'B-PRO', 5: 'B-NAME', 6: 'B-EDU', 7: 'B-LOC',
               8: 'B-TITLE', 9: 'I-CONT', 10: 'I-ORG', 11: 'I-RACE', 12: 'I-PRO', 13: 'I-NAME', 14: 'I-EDU',
               15: 'I-LOC', 16: 'I-TITLE'}

    sent = '小明，中国人，在清华大学就读计算机专业。'
    inputs = tokenizer(sent, return_tensors='pt')
    outputs, t = model((inputs.input_ids, inputs.attention_mask))
    sns.set(font_scale=0.6)
    sns.heatmap(t, annot=True,  linewidths=0.005)
    plt.show()
    for output in outputs:
        res = [tag_map[i] for i in output[1:-1]]
        print(res)
        print(get_entities(res))
        entities = []
        for entity in get_entities(res):
            entities.append((entity[0], sent[entity[1]: entity[2] + 1]))
        print(entities)


if __name__ == '__main__':
    main()
